import torch
import os
import pickle  # 用于保存和加载特征
from tqdm import tqdm  # 用于显示进度条
from omni_model.omni_space import OmniBind_Base  # 导入模型类
from topk_valor import load_features_from_disk, save_features_to_disk, create_batches_from_json  # 从 topk_valor.py 导入相关函数

def compute_similarity_matrices(
    model: OmniBind_Base,
    primary_modality1: str,
    primary_modality2: str,
    intermediate_modality: str,
    json_path: str,
    base_folder: str,
    audio_folder: str,
    batch_size: int = 32
) -> tuple:
    """
    计算主模态1 -> 中间模态和主模态2 -> 中间模态的相似度矩阵。

    Args:
        model (OmniBind_Base): 模型实例。
        primary_modality1 (str): 主模态1名称。
        primary_modality2 (str): 主模态2名称。
        intermediate_modality (str): 中间模态名称。
        json_path (str): JSON 文件路径。
        base_folder (str): 根文件夹路径。
        audio_folder (str): 音频文件夹路径。
        batch_size (int): 批次大小。

    Returns:
        tuple: 包含 feat1_mid 和 feat2_mid 的相似度矩阵。
    """
    # 定义保存/加载特征的文件路径
    feat1_path = f"{primary_modality1}_features.pkl"
    feat_mid_path = f"{intermediate_modality}_features.pkl"
    feat2_path = f"{primary_modality2}_features.pkl"

    # 加载或提取主模态1的特征
    if os.path.exists(feat1_path):
        print(f"Loading {primary_modality1} features from disk...")
        feat1 = load_features_from_disk(feat1_path)
    else:
        print(f"Extracting {primary_modality1} features...")
        batches1 = create_batches_from_json(json_path, base_folder, batch_size, primary_modality1, audio_folder)
        feat1_list = []
        for batch in batches1:
            if primary_modality1 == "audio":
                feat1_batch = model.emb_audios(batch)
            elif primary_modality1 == "image":
                feat1_batch = model.emb_images(batch)
            else:
                feat1_batch = model.emb_texts(batch)
            feat1_list.append(feat1_batch)
        feat1 = torch.cat(feat1_list, dim=0)
        save_features_to_disk(feat1, feat1_path)

    # 加载或提取中间模态的特征
    if os.path.exists(feat_mid_path):
        print(f"Loading {intermediate_modality} features from disk...")
        feat_mid = load_features_from_disk(feat_mid_path)
    else:
        print(f"Extracting {intermediate_modality} features...")
        batches_mid = create_batches_from_json(json_path, base_folder, batch_size, intermediate_modality, audio_folder)
        feat_mid_list = []
        for batch in batches_mid:
            if intermediate_modality == "audio":
                feat_mid_batch = model.emb_audios(batch)
            elif intermediate_modality == "image":
                feat_mid_batch = model.emb_images(batch)
            else:
                feat_mid_batch = model.emb_texts(batch)
            feat_mid_list.append(feat_mid_batch)
        feat_mid = torch.cat(feat_mid_list, dim=0)
        save_features_to_disk(feat_mid, feat_mid_path)

    # 加载或提取主模态2的特征
    if os.path.exists(feat2_path):
        print(f"Loading {primary_modality2} features from disk...")
        feat2 = load_features_from_disk(feat2_path)
    else:
        print(f"Extracting {primary_modality2} features...")
        batches2 = create_batches_from_json(json_path, base_folder, batch_size, primary_modality2, audio_folder)
        feat2_list = []
        for batch in batches2:
            if primary_modality2 == "audio":
                feat2_batch = model.emb_audios(batch)
            elif primary_modality2 == "image":
                feat2_batch = model.emb_images(batch)
            else:
                feat2_batch = model.emb_texts(batch)
            feat2_list.append(feat2_batch)
        feat2 = torch.cat(feat2_list, dim=0)
        save_features_to_disk(feat2, feat2_path)

    # 打印模态1数量
    print(f"Number of {primary_modality1} features: {feat1.size(0)}")

    # 计算主模态1 -> 中间模态的相似度矩阵
    print(f"Computing {primary_modality1} -> {intermediate_modality} similarity matrix...")
    feat1_mid = feat1 @ feat_mid.T

    # 计算主模态2 -> 中间模态的相似度矩阵
    print(f"Computing {primary_modality2} -> {intermediate_modality} similarity matrix...")
    feat2_mid = feat2 @ feat_mid.T

    return feat1_mid, feat2_mid, feat_mid


def save_intersection_of_r1_features(
    feat1_mid: torch.Tensor,
    feat2_mid: torch.Tensor,
    feat_mid: torch.Tensor,
    output_filepath: str,
    topk: int = 1  # 新增参数 topk，默认为 1
):
    """
    保存 feat1 -> feat_mid 和 feat2 -> feat_mid 的 r1 中的 feat_mid 特征交集到文件。

    Args:
        feat1_mid (torch.Tensor): feat1 -> feat_mid 的相似度矩阵。
        feat2_mid (torch.Tensor): feat2 -> feat_mid 的相似度矩阵。
        feat_mid (torch.Tensor): 中间模态的特征矩阵。
        output_filepath (str): 输出文件路径。
        topk (int): 选择前 k 个相似特征。
    """
    # 计算 feat1 -> feat_mid 的 topk 中的 feat_mid 索引
    print("Computing topk indices for feat1 -> feat_mid...")
    feat1_r1_indices = []
    for i in range(feat1_mid.size(0)):
        top_indices = torch.topk(feat1_mid[i], topk, dim=0).indices.tolist()
        feat1_r1_indices.extend(top_indices)

    # 计算 feat2 -> feat_mid 的 topk 中的 feat_mid 索引
    print("Computing topk indices for feat2 -> feat_mid...")
    feat2_r1_indices = []
    for i in range(feat2_mid.size(0)):
        top_indices = torch.topk(feat2_mid[i], topk, dim=0).indices.tolist()
        feat2_r1_indices.extend(top_indices)

    # 计算交集
    print("Computing intersection of topk indices...")
    feat1_r1_set = set(feat1_r1_indices)
    feat2_r1_set = set(feat2_r1_indices)
    intersection_indices = list(feat1_r1_set.intersection(feat2_r1_set))

    # 打印交集数量
    print(f"Number of features in the intersection: {len(intersection_indices)}")

    # 提取交集特征
    print("Extracting intersection features...")
    intersection_features = feat_mid[intersection_indices]

    # 保存交集特征到文件
    print(f"Saving intersection features to {output_filepath}...")
    with open(output_filepath, 'wb') as f:
        pickle.dump(intersection_features, f)

    print("Intersection features saved successfully.")


def save_same_id_r1_features(
    feat1_mid: torch.Tensor,
    feat2_mid: torch.Tensor,
    feat_mid: torch.Tensor,
    output_filepath: str
):
    """
    保存 feat1 和 feat2 的 r1 中具有相同索引的 feat_mid 特征到文件。

    Args:
        feat1_mid (torch.Tensor): feat1 -> feat_mid 的相似度矩阵。
        feat2_mid (torch.Tensor): feat2 -> feat_mid 的相似度矩阵。
        feat_mid (torch.Tensor): 中间模态的特征矩阵。
        output_filepath (str): 输出文件路径。
    """
    # 计算 feat1 -> feat_mid 的 r1 中的 feat_mid 索引
    print("Computing r1 indices for feat1 -> feat_mid...")
    feat1_r1_indices = []
    for i in range(feat1_mid.size(0)):
        top_index = torch.topk(feat1_mid[i], 1, dim=0).indices.item()
        feat1_r1_indices.append(top_index)

    # 计算 feat2 -> feat_mid 的 r1 中的 feat_mid 索引
    print("Computing r1 indices for feat2 -> feat_mid...")
    feat2_r1_indices = []
    for i in range(feat2_mid.size(0)):
        top_index = torch.topk(feat2_mid[i], 1, dim=0).indices.item()
        feat2_r1_indices.append(top_index)

    # 找到 feat1 和 feat2 的 r1 中具有相同索引的特征
    print("Finding same id r1 indices...")
    same_id_indices = [i for i, (idx1, idx2) in enumerate(zip(feat1_r1_indices, feat2_r1_indices)) if idx1 == idx2]

    # 打印相同 id 的数量
    print(f"Number of features with the same r1 indices: {len(same_id_indices)}")

    # 提取相同 id 的特征
    print("Extracting same id features...")
    same_id_features = feat_mid[same_id_indices]

    # 保存相同 id 特征到文件
    print(f"Saving same id features to {output_filepath}...")
    with open(output_filepath, 'wb') as f:
        pickle.dump(same_id_features, f)

    print("Same id features saved successfully.")


def save_matching_r1_features(
    feat1_mid: torch.Tensor,
    feat2_mid: torch.Tensor,
    feat_mid: torch.Tensor,
    output_filepath: str
):
    """
    保存同时符合 feat1 和 feat2 的 r1 召回且 ID 分别对应的 feat_mid 特征到文件。

    Args:
        feat1_mid (torch.Tensor): feat1 -> feat_mid 的相似度矩阵。
        feat2_mid (torch.Tensor): feat2 -> feat_mid 的相似度矩阵。
        feat_mid (torch.Tensor): 中间模态的特征矩阵。
        output_filepath (str): 输出文件路径。
    """
    # 计算 feat1 -> feat_mid 的 r1 中的 feat_mid 索引
    print("Computing r1 indices for feat1 -> feat_mid...")
    feat1_r1_indices = []
    for i in range(feat1_mid.size(0)):
        top_index = torch.topk(feat1_mid[i], 1, dim=0).indices.item()
        feat1_r1_indices.append(top_index)

    # 计算 feat2 -> feat_mid 的 r1 中的 feat_mid 索引
    print("Computing r1 indices for feat2 -> feat_mid...")
    feat2_r1_indices = []
    for i in range(feat2_mid.size(0)):
        top_index = torch.topk(feat2_mid[i], 1, dim=0).indices.item()
        feat2_r1_indices.append(top_index)

    # 找到同时符合 feat1 和 feat2 的 r1 召回且 ID 分别对应的特征
    print("Finding features recalled by both r1 with corresponding IDs...")
    corresponding_indices = [
        idx for idx, (idx1, idx2) in enumerate(zip(feat1_r1_indices, feat2_r1_indices)) if idx1 == idx and idx2 == idx
    ]

    # 打印符合条件的特征数量
    print(f"Number of features recalled by both r1 with corresponding IDs: {len(corresponding_indices)}")

    # 提取符合条件的特征
    print("Extracting features recalled by both r1 with corresponding IDs...")
    corresponding_features = feat_mid[corresponding_indices]

    # 保存特征到文件
    print(f"Saving features recalled by both r1 with corresponding IDs to {output_filepath}...")
    with open(output_filepath, 'wb') as f:
        pickle.dump(corresponding_features, f)

    print("Features recalled by both r1 with corresponding IDs saved successfully.")


# 示例用法
if __name__ == "__main__":
    # 假设模型已加载
    # model = OmniBind_Base(pretrained=True).cuda().eval()

    # 输入参数
    json_path = "/data/jzw/valor-32k-annotations/desc_val_filtered.json"
    base_folder = "/data/jzw/processed_valor32k_val/frames"
    audio_folder = "/data/jzw/processed_valor32k_val/audio"
    primary_modality1 = "audio"
    primary_modality2 = "text"
    intermediate_modality = "image"
    batch_size = 32

    # 调用函数计算相似度矩阵
    feat1_mid, feat2_mid, feat_mid = compute_similarity_matrices(
        model=None,
        primary_modality1=primary_modality1,
        primary_modality2=primary_modality2,
        intermediate_modality=intermediate_modality,
        json_path=json_path,
        base_folder=base_folder,
        audio_folder=audio_folder,
        batch_size=batch_size
    )

    topk = 5
    output_filepath = f"intersection_features_top{topk}.pkl"
    # 保存交集特征
    save_intersection_of_r1_features(
        feat1_mid=feat1_mid,
        feat2_mid=feat2_mid,
        feat_mid=feat_mid,
        output_filepath=output_filepath,
        topk=topk  # 选择前 k 个相似特征
    )

    # # 保存相同 id 特征
    # same_id_output_filepath = "same_id_features.pkl"
    # save_same_id_r1_features(
    #     feat1_mid=feat1_mid,
    #     feat2_mid=feat2_mid,
    #     feat_mid=feat_mid,
    #     output_filepath=same_id_output_filepath
    # )

    # # 保存同时符合两个模态 r1 召回且 ID 分别对应的特征
    # matching_r1_output_filepath = "matching_r1_features.pkl"
    # save_matching_r1_features(
    #     feat1_mid=feat1_mid,
    #     feat2_mid=feat2_mid,
    #     feat_mid=feat_mid,
    #     output_filepath=matching_r1_output_filepath
    # )